import pandas as pd
from service.das.stock_day_das import StockDateDayDas
from service.das.stock_kline_das import KLineData, StockKLineDas


class StockAnalyseService:

    def calculate_and_store_kline(self, stock_code: str):
    """
    计算并存储指定股票的周K、月K、年K数据。
    :param stock_code: 股票代码。
    """
    # 从 DAS 获取所有日K数据
    stockDayDas = StockDateDayDas()
    day_data = stockDayDas.search_by_day_all(stock_code)
    
    if not day_data:
        print(f"No day K-line data available for stock_code={stock_code}.")
        return

    # 转换为 DataFrame
    data = pd.DataFrame([day.__dict__ for day in day_data])
    data["trade_date"] = pd.to_datetime(data["trade_date"])

    # 定义需要生成的 K 线类型及分组规则
    kline_types = {
        "W": data["trade_date"].dt.to_period("W").apply(lambda r: r.start_time),
        "M": data["trade_date"].dt.to_period("M").apply(lambda r: r.start_time),
        "Y": data["trade_date"].dt.to_period("Y").apply(lambda r: r.start_time),
    }

    # 初始化 KLine DAS
    klineDas = StockKLineDas()

    # 清除该股票的历史 K 线数据
    klineDas.delete_kline_data(stock_code)

    # 逐个生成并存储每种类型的 K 线数据
    for kline_type, group_data in kline_types.items():
        data["group"] = group_data

        # 聚合计算
        grouped = data.groupby(["stock_code", "group"])
        aggregated = grouped.agg(
            open=("open", "first"),
            high=("high", "max"),
            low=("low", "min"),
            close=("close", "last"),
            vol=("vol", "sum"),
            amount=("amount", "sum"),
        ).reset_index()

        # 存储聚合后的数据
        klineDas.store_aggregated_kline(aggregated, kline_type)

    print(f"Successfully calculated and stored K-line data for stock_code={stock_code}.")
